from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
import MyDataloader
from CNN import CNN
import torch
import torchvision
import cv2
import numpy as np
import matplotlib as mpl

device = torch.device("cuda")
images_folder_test_x = r"test1_x_1" + "/"
images_folder_test_y = r"test1_y_1" + "/"


def rgb2hex(r, g, b):
    """将 RGB 颜色转换为十六进制颜色"""
    hex_color = "#{:02x}{:02x}{:02x}".format(r, g, b)
    return hex_color


def testImage():
    test_data = MyDataloader.MyDataset(95, images_folder_test_x, images_folder_test_y,
                                       transform=torchvision.transforms.ToTensor())
    test_loader = DataLoader(dataset=test_data, batch_size=12, shuffle=False, drop_last=True)

    model = CNN()
    model.load_state_dict(torch.load("CNN_NC.pth"))

    for i, data in enumerate(test_loader):
        input = data[0].cuda()
        target = data[1].cuda()
        input2 = input.reshape(12, 448, 304)
        prediction = model(input2)

        # 下面进行预测图片的生成
        prediction = prediction.cpu()
        prediction = prediction.detach().numpy().reshape(448, 304)
        mycolors = [rgb2hex(0, 6, 134), rgb2hex(0, 105, 218), rgb2hex(0, 122, 233),
                    rgb2hex(0, 142, 249), rgb2hex(13, 159, 255), rgb2hex(61, 179, 255), rgb2hex(110, 198, 255),
                    rgb2hex(154, 215, 255),
                    rgb2hex(199, 233, 255), rgb2hex(247, 252, 255)]
        bounds = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]
        cmap = mpl.colors.ListedColormap(mycolors)
        norm = mpl.colors.BoundaryNorm(bounds, cmap.N)

        fig = plt.figure(figsize=(3.04, 4.48), dpi=130)
        plt.imshow(prediction, cmap=cmap, norm=norm)
        # fig.colorbar()
        plt.axis('off')
        fig.savefig('CNN_NC/{}pred.png'.format(i + 1), bbox_inches='tight', pad_inches=0)
        plt.close()

        image1 = cv2.imread('1.png')
        image2 = cv2.imread('CNN_NC/{}pred.png'.format(i + 1))
        gray_color = (119, 119, 119)
        gray_color = np.array(gray_color)
        gray_coords = np.where(np.all(image1 == gray_color, axis=-1))
        image2[gray_coords[0], gray_coords[1]] = gray_color
        cv2.imwrite('CNN_NC/{}pred.png'.format(i + 1), image2)

        target = target[11]
        target = target.cpu()
        target = target.detach().numpy().reshape(448, 304)
        fig2 = plt.figure(figsize=(3.04, 4.48), dpi=130)
        plt.imshow(target, cmap=cmap, norm=norm)
        # fig.colorbar()
        plt.axis('off')
        fig2.savefig('CNN_NC/{}true.png'.format(i + 1), bbox_inches='tight', pad_inches=0)
        plt.close()

        image3 = cv2.imread('CNN_NC/{}true.png'.format(i + 1))
        gray_color = (119, 119, 119)
        gray_color = np.array(gray_color)
        gray_coords = np.where(np.all(image1 == gray_color, axis=-1))
        image3[gray_coords[0], gray_coords[1]] = gray_color
        cv2.imwrite('CNN_NC/{}true.png'.format(i + 1), image3)

        # 存npy
        prediction = np.round(prediction, decimals=4)
        np.save('CNN_npy/{}pred.npy'.format(i + 1), prediction)

        target = np.round(target, decimals=4)
        np.save('CNN_npy/{}true.npy'.format(i + 1), target)


        del input, target, prediction


if __name__ == '__main__':
    testImage()
